package scu.maqiang.nla;

import scu.maqiang.fes.BVPType;
import scu.maqiang.fes.FES1L21;
import scu.maqiang.mesh.Mesh1L2;
import scu.maqiang.mesh.ScalarFunc;
import scu.maqiang.mesh.Tecplot;
import scu.maqiang.numeric.Direct;
import scu.maqiang.numeric.MVO;
import scu.maqiang.numeric.NewIterSSolver;
import scu.maqiang.numeric.SRMatrix;

import java.util.Arrays;

public class NonLinearPossion1D {
    public static void main(String[] args) {
        //DirectIterationMethod();
        //NewtonIterationMethod();
        //DirectIterationMethod_HeatTransfer();
        NewtonIterationMethod_HeatTransfer();
        //DirectIterationMethod_example452();
        //NewtonIteration_LargeDeformation();
        //NewtonIteration_MaterialNonlinearity();
        //DirectIterationMethod_LargeDeformation();
    }

    public static void DirectIterationMethod() {
        Mesh1L2 mesh = new Mesh1L2();
        mesh.line(0, 1, 4);
        mesh.displayNodes();
        mesh.displayElements();
        mesh.displayBoundarys();

        FES1L21 fs = new FES1L21(mesh);
        SRMatrix A = new SRMatrix(fs.GetNdof());
        double[] RHS = new double[fs.GetNdof()];
        double[] x = {1.0, 1.0, 1.0, 1.0, Math.sqrt(2)};

        fs.assembleSource(new double[]{-1.0}, BVPType.COMMON, RHS);
        double[] dx = new double[fs.GetNdof()];

        NewIterSSolver solver = new NewIterSSolver();
        int maxIteration = 100000;
        int k = 0;

        double[] zeroVec = new double[fs.GetNdof()];
        double[] a = new double[fs.GetNdof()];
        Arrays.setAll(a, i -> 1.0);
        double[][] coef = {zeroVec, zeroVec, zeroVec, zeroVec, a, zeroVec, zeroVec, zeroVec, zeroVec, zeroVec, x};
        double[] constCoef = {0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0};
        double[][] vecCoef = {x};
        while(k < maxIteration) {
            fs.assembleStiff(coef, BVPType.NONLINEAR_1, A);
            //fs.assembleStiff(constCoef, vecCoef, BVPType.NONLINEAR_1, A);
            fs.applyBC_MBN(A, RHS, Direct.All, Math.sqrt(2), 2);
            solver.setMatrix(A);
            solver.PCGSSOR(RHS, dx, 1.5, 1);
            double[] diff = MVO.add(dx, -1.0, x);
            if(MVO.L2Norm(diff) <= 1.0e-5 * MVO.L2Norm(x)) {
                break;
            }
            Arrays.setAll(x, i -> dx[i]);
            System.out.println(Arrays.toString(x));
            A.clear();
            k++;
        }
    }

    public static void DirectIterationMethod_HeatTransfer() {
        Mesh1L2 mesh = new Mesh1L2();
        mesh.line(0, 0.18, 8);
        mesh.displayNodes();
        mesh.displayElements();
        mesh.displayBoundarys();

        FES1L21 fs = new FES1L21(mesh);
        SRMatrix A = new SRMatrix(fs.GetNdof());
        double[] RHS = new double[fs.GetNdof()];
        double[] x = new double[fs.GetNdof()];
        x[0] = 500;
        x[fs.GetNdof() - 1] = 300;


        double[] dx = new double[fs.GetNdof()];
        NewIterSSolver solver = new NewIterSSolver();
        int maxIteration = 100000;
        int k = 0;

        double[] zeroVec = new double[fs.GetNdof()];

        double k0 = 0.2;
        double k1 = 2e-3;

        double[] k0Vec = new double[fs.GetNdof()];
        Arrays.setAll(k0Vec, i-> k0);
        double[] k0k1Vec = new double[fs.GetNdof()];
        Arrays.setAll(k0k1Vec, i-> k0 * k1);
        double[][] coef = {k0Vec, zeroVec, zeroVec, zeroVec, k0k1Vec, zeroVec, zeroVec, zeroVec, zeroVec, zeroVec, x};
        while(k < maxIteration) {
            fs.assembleStiff(coef, BVPType.NONLINEAR_1, A);
            fs.applyBC_MBN(A, RHS, Direct.All, 500, 1);
            fs.applyBC_MBN(A, RHS, Direct.All, 300, 2);
            solver.setMatrix(A);
            solver.PCGSSOR(RHS, dx, 1.5, 1);
            double[] diff = MVO.add(dx, -1.0, x);
            if(MVO.L2Norm(diff) <= 1.0e-5 * MVO.L2Norm(x)) {
                break;
            }
            Arrays.setAll(x, i -> dx[i]);
            System.out.println(Arrays.toString(x));
            A.clear();
            k++;
        }
    }

    /**
     *  @brief 求解 -d^2 u/dx^2 + 2 u^3 = 0, 1 < x < 2,
     *  s.t. u(1) = 1, du/dx + u^2|_{x = 2} = 0
     */
    public static void DirectIterationMethod_example452() {
        Mesh1L2 mesh = new Mesh1L2();
        mesh.line(1, 2, 8);
        mesh.displayNodes();
        mesh.displayElements();
        mesh.displayBoundarys();

        FES1L21 fs = new FES1L21(mesh);
        SRMatrix A = new SRMatrix(fs.GetNdof());
        double[] RHS = new double[fs.GetNdof()];
        double[] x = new double[fs.GetNdof()];
        Arrays.setAll(x, i-> 0.5);
        x[0] = 1.0;

//        fs.assembleStiff(new double[]{1.0}, BVPType.COMMON, A);

        double[] a = new double[fs.GetNdof()];
        Arrays.setAll(a, i -> 2.0);
        double[] zeroVec = new double[fs.GetNdof()];
        double[][] coef = {zeroVec, zeroVec, zeroVec, zeroVec, zeroVec, a, zeroVec, zeroVec, zeroVec, zeroVec, x};

//        double[] coef = new double[fs.GetNdof()];
//        Arrays.setAll(coef, i-> 2 * x[i] * x[i]);
//        fs.assembleHeatMass(coef, BVPType.NONLINEAR_1, A);
//        fs.assembleHeatConvectSurface(new double[][]{x}, BVPType.CONSIST_MATRIX, A, 2);
        //A.addElement(fs.GetNdof() - 1, fs.GetNdof() - 1, x[fs.GetNdof() - 1]);
        double[] dx = new double[fs.GetNdof()];
        NewIterSSolver solver = new NewIterSSolver();
        int maxIteration = 100000;
        int k = 0;
        while(k < maxIteration) {
            A.clear();
            fs.assembleStiff(new double[]{1.0}, BVPType.COMMON, A);
            fs.assembleMass(coef, BVPType.NONLINEAR_1, A);
            fs.assembleBoundaryMass(new double[][]{x}, BVPType.CONSIST_MASS, A, 2);

            fs.applyBC_MBN(A, RHS, Direct.All, 1, 1);
            solver.setMatrix(A);
            solver.PCGSSOR(RHS, dx, 1.5, 1);
            double[] diff = MVO.add(dx, -1.0, x);
            if(MVO.L2Norm(diff) < 1.0e-10) {
                break;
            }
            Arrays.setAll(x, i -> dx[i]);
            System.out.println(Arrays.toString(x));

            //A.addElement(fs.GetNdof() - 1, fs.GetNdof() - 1, x[fs.GetNdof() - 1]);
            A.clear();
            k++;
        }
        ScalarFunc exactFunc = (xy, label, param) -> 1.0/xy[0];
        double[] exactU = fs.valueFromFunc(exactFunc, null);
        System.out.println(MVO.distance(exactU, x));
    }

    public static void NewtonIterationMethod() {
        Mesh1L2 mesh = new Mesh1L2();
        mesh.line(0, 1, 4);
        FES1L21 fs = new FES1L21(mesh);
        double[] RHS = new double[fs.GetNdof()];
        double[] x = {1.0, 1.0, 1.0, 1.0, Math.sqrt(2)};
        double[] k = new double[fs.GetNdof()];
        double[] a = new double[fs.GetNdof()];
        Arrays.setAll(a, i->1.0);
        double[] b = new double[fs.GetNdof()];
        fs.assembleSource(new double[]{-1.0}, BVPType.COMMON, RHS);

        SRMatrix A0 = new SRMatrix(fs.GetNdof());
        SRMatrix A = new SRMatrix(fs.GetNdof());
        double[] dx = new double[fs.GetNdof()];
        double[] zeroVec = new double[fs.GetNdof()];
        double[][] coef = {zeroVec, zeroVec, zeroVec, zeroVec, a, zeroVec, zeroVec, zeroVec, zeroVec, zeroVec, x};
        int maxIteration = 10000;
        int iter = 0;
        NewIterSSolver solver = new NewIterSSolver();
        while (iter < maxIteration) {
            fs.assembleStiff(coef, BVPType.NONLINEAR_1, A0);
            double[] R = MVO.add(RHS, -1.0, A0.mul(x));
            A0.clear();
            fs.assembleStiff(coef, BVPType.NONLINEAR_2, A);
            fs.applyBC_MBN(A, RHS, Direct.All, 0.0, 2);
            solver.setMatrix(A);
            solver.PCGSSSOR(R, dx, 1.5, 1);
            Arrays.setAll(x, i -> x[i] + dx[i]);
            if(MVO.L2Norm(dx) < 1.0e-10 && MVO.L2Norm(R) < 1.0e-10) {
                break;
            }
            A.clear();
            iter++;
        }

        System.out.println(Arrays.toString(dx));
        System.out.println(Arrays.toString(x));

        //System.out.println(Arrays.toString(R));


        //A.sort();
        //System.out.println(A.display());

    }

    public static void NewtonIterationMethod_HeatTransfer() {
        Mesh1L2 mesh = new Mesh1L2();
        mesh.line(0, 1, 8);
        FES1L21 fs = new FES1L21(mesh);

        double[] RHS = new double[fs.GetNdof()];
        double[] x = {500, 475, 450, 425, 400, 375, 350, 325, 300};

        double[] b = new double[fs.GetNdof()];

        SRMatrix A0 = new SRMatrix(fs.GetNdof());
        SRMatrix A = new SRMatrix(fs.GetNdof());
        double[] dx = new double[fs.GetNdof()];
        double[] myCoef = new double[fs.GetNdof()];

        double[] k0Vec = new double[fs.GetNdof()];
        Arrays.setAll(k0Vec, i-> 0.2);
        double[] k0k1Vec = new double[fs.GetNdof()];
        Arrays.setAll(k0k1Vec, i-> 0.2 * 2e-3);
        double[] zeroVec = new double[fs.GetNdof()];
        double[][] coef = {k0Vec, zeroVec, zeroVec, zeroVec, k0k1Vec, zeroVec, zeroVec, zeroVec, zeroVec, zeroVec, x};
        //double[][] coef = {k0, k1, zeoVec, x};

        int maxIteration = 10000;
        int iter = 0;
        NewIterSSolver solver = new NewIterSSolver();
        while (iter < maxIteration) {
            //1. 计算余量
            A0.clear();
            fs.assembleStiff(coef, BVPType.NONLINEAR_1, A0);
            double[] R = MVO.add(RHS, -1.0, A0.mul(x));
            //2. 组装切线矩阵
            A.clear();
            fs.assembleStiff(coef, BVPType.NONLINEAR_2, A);

            //3. 施加边界条件
            fs.applyBC_MBN(A, RHS, Direct.All, 0.0, 1, 2);

            //4. 求解增量dx
            solver.setMatrix(A);
            solver.PCGSSSOR(R, dx, 1.5, 1);

            //5. 计算整体解
            Arrays.setAll(x, i -> x[i] + dx[i]);

            //6. 检验收敛性
            if(MVO.L2Norm(dx) < 1.0e-10 * MVO.L2Norm(x) && MVO.L2Norm(R) < 1.0e-10 * MVO.L2Norm(RHS)) {
                break;
            }
            iter++;
        }

        System.out.println(Arrays.toString(dx));
        System.out.println(Arrays.toString(x));

        //System.out.println(Arrays.toString(R));


        //A.sort();
        //System.out.println(A.display());

    }

    public static void NewtonIterationMethod_Example452() {
        Mesh1L2 mesh = new Mesh1L2();
        mesh.line(1, 2, 4);

        FES1L21 fs = new FES1L21(mesh);

        double[] RHS = new double[fs.GetNdof()];
        double[] x = new double[fs.GetNdof()];
        Arrays.setAll(x, i-> 0.5);
        x[0] = 1.0;

        double[] b = new double[fs.GetNdof()];

        SRMatrix A0 = new SRMatrix(fs.GetNdof());
        SRMatrix A = new SRMatrix(fs.GetNdof());
        double[] dx = new double[fs.GetNdof()];
        double[] myCoef = new double[fs.GetNdof()];

//        double[] k0 = new double[fs.GetNdof()];
//        Arrays.setAll(k0, i-> 0.2);
//        double[] k1 = new double[fs.GetNdof()];
//        Arrays.setAll(k1, i-> 0.2 * 2e-3);
//        double[] zeoVec = new double[fs.GetNdof()];
//        //double[][] coef = {k0, zeoVec, zeoVec, k1, zeoVec, zeoVec, zeoVec, x};
//        double[][] coef = {k0, k1, zeoVec, x};

        double[] a = new double[fs.GetNdof()];
        Arrays.setAll(a, i -> 2.0);
        double[] zeroVec = new double[fs.GetNdof()];
        double[][] coef = {zeroVec, zeroVec, zeroVec, zeroVec, zeroVec, a, zeroVec, zeroVec, zeroVec, zeroVec, x};

        int maxIteration = 10000;
        int iter = 0;
        NewIterSSolver solver = new NewIterSSolver();
        while (iter < maxIteration) {
            //1. 计算余量
            //Arrays.setAll(myCoef, i -> k0[i]  + k1[i] * x[i]);
            fs.assembleStiff(new double[][]{myCoef}, BVPType.COMMON, A0);
            double[] R = MVO.add(RHS, -1.0, A0.mul(x));
            A0.clear();
            //2. 组装切线矩阵
            fs.assembleStiff(coef, BVPType.NONLINEAR_1, A);

            //3. 施加边界条件
            fs.applyBC_MBN(A, RHS, Direct.All, 0.0, 1, 2);

            //4. 求解增量dx
            solver.setMatrix(A);
            solver.PCGSSSOR(R, dx, 1.5, 1);

            //5. 计算整体解
            Arrays.setAll(x, i -> x[i] + dx[i]);

            //6. 检验收敛性
            if(MVO.L2Norm(dx) < 1.0e-10 * MVO.L2Norm(x) && MVO.L2Norm(R) < 1.0e-10 * MVO.L2Norm(RHS)) {
                break;
            }
            iter++;
        }

        System.out.println(Arrays.toString(dx));
        System.out.println(Arrays.toString(x));

        //System.out.println(Arrays.toString(R));


        //A.sort();
        //System.out.println(A.display());

    }

    public static void DirectIterationMethod_LargeDeformation() {
        Mesh1L2 mesh = new Mesh1L2();
        mesh.line(0, 1, 8);
        mesh.displayNodes();
        mesh.displayElements();
        mesh.displayBoundarys();

        FES1L21 fs = new FES1L21(mesh);
        SRMatrix A = new SRMatrix(fs.GetNdof());
        double[] RHS = new double[fs.GetNdof()];
        double[] x = new double[fs.GetNdof()];

        //fs.assembleHeatSource(new double[]{-1.0}, BVPType.COMMON, RHS);
        double[] dx = new double[fs.GetNdof()];

        NewIterSSolver solver = new NewIterSSolver();
        int maxIteration = 100000;
        int k = 0;

        double[] coef0 = new double[fs.GetNdof()];
        Arrays.setAll(coef0, i-> 1.0);
        double[] coefdU = new double[fs.GetNdof()];
        Arrays.setAll(coefdU, i-> 1.5);
        double[] coefdU2 = new double[fs.GetNdof()];
        Arrays.setAll(coefdU2, i ->0.5);
        double[] zeroVec = new double[fs.GetNdof()];
        double[][] coef = {coef0, zeroVec, zeroVec, zeroVec, zeroVec, zeroVec, zeroVec, coefdU, coefdU2, zeroVec, x};
        fs.assembleFlux(new double[]{0.2}, BVPType.COMMON, RHS, 2);
        while(k < maxIteration) {
            fs.assembleStiff(coef, BVPType.NONLINEAR_1, A);
            fs.applyBC_MBN(A, RHS, Direct.All, 0.0, 1);
            solver.setMatrix(A);
            solver.PCGSSOR(RHS, dx, 1.5, 1);
            double[] diff = MVO.add(dx, -1.0, x);
            if(MVO.L2Norm(diff) <= 1.0e-5 * MVO.L2Norm(x)) {
                break;
            }
            Arrays.setAll(x, i -> dx[i]);
            System.out.println(Arrays.toString(x));
            A.clear();
            k++;
        }
    }

    public static void NewtonIteration_LargeDeformation() {
        Mesh1L2 mesh = new Mesh1L2();
        mesh.line(0, 1, 8);
        FES1L21 fs = new FES1L21(mesh);

        double[] RHS = new double[fs.GetNdof()];
        double[] x = new double[fs.GetNdof()];
        double[] b = new double[fs.GetNdof()];

        SRMatrix A0 = new SRMatrix(fs.GetNdof());
        SRMatrix A = new SRMatrix(fs.GetNdof());
        double[] dx = new double[fs.GetNdof()];

        double[] coef0 = new double[fs.GetNdof()];
        Arrays.setAll(coef0, i-> 1.0);
        double[] coefdU = new double[fs.GetNdof()];
        Arrays.setAll(coefdU, i-> 1.5);
        double[] coefdU2 = new double[fs.GetNdof()];
        Arrays.setAll(coefdU2, i ->0.5);
        double[] zeroVec = new double[fs.GetNdof()];
        double[][] coef = {coef0, zeroVec, zeroVec, zeroVec, zeroVec, zeroVec, zeroVec, coefdU, coefdU2, zeroVec, x};
        //double[][] coef = {k0, k1, zeoVec, x};

        double[] disp = new double[26];
        double[] force = new double[26];
        double presentforce = 0.0;
        //fs.assembleHeatFlux(new double[]{5.0}, BVPType.COMMON, RHS, 2);
        int maxIteration = 10000;
        int iter = 0;
        NewIterSSolver solver = new NewIterSSolver();
        int idx = 1;
        while(idx < 26) {
            presentforce += 0.2;
            fs.assembleFlux(new double[]{0.2}, BVPType.COMMON, RHS, 2);
            iter = 0;
            while (iter < maxIteration) {
                //1. 计算余量
                A0.clear();
                fs.assembleStiff(coef, BVPType.NONLINEAR_1, A0);
                double[] R = MVO.add(RHS, -1.0, A0.mul(x));
                //2. 组装切线矩阵
                A.clear();
                fs.assembleStiff(coef, BVPType.NONLINEAR_2, A);

                //3. 施加边界条件
                fs.applyBC_MBN(A, RHS, Direct.All, 0.0, 1);

                //4. 求解增量dx
                solver.setMatrix(A);
                solver.PCGSSSOR(R, dx, 1.5, 0);

                //5. 计算整体解
                Arrays.setAll(x, i -> x[i] + dx[i]);

                //6. 检验收敛性
                if(MVO.L2Norm(dx) < 1.0e-10 * MVO.L2Norm(x) && MVO.L2Norm(R) < 1.0e-10 * MVO.L2Norm(RHS)) {
                    break;
                }
                iter++;
            }
            force[idx] = presentforce;
            disp[idx] = x[fs.GetNdof() - 1];
            System.out.println(force[idx] + "\t\t" + disp[idx]);
            idx++;
        }

        Tecplot.LineXY("LoadVsDisp.dat", force, new double[][]{disp});
//        System.out.println(Arrays.toString(dx));
//        System.out.println(Arrays.toString(x));

    }

    public static void NewtonIteration_MaterialNonlinearity() {
        Mesh1L2 mesh = new Mesh1L2();
        mesh.line(0, 1, 8);
        FES1L21 fs = new FES1L21(mesh);

        double[] RHS = new double[fs.GetNdof()];
        double[] x = new double[fs.GetNdof()];
        double[] b = new double[fs.GetNdof()];

        SRMatrix A0 = new SRMatrix(fs.GetNdof());
        SRMatrix A = new SRMatrix(fs.GetNdof());
        double[] dx = new double[fs.GetNdof()];

        double[] coef0 = new double[fs.GetNdof()];
        Arrays.setAll(coef0, i-> 1.0);
        double[] coefdU = new double[fs.GetNdof()];
        Arrays.setAll(coefdU, i-> -0.2);
        double[] zeroVec = new double[fs.GetNdof()];
        double[][] coef = {coef0, zeroVec, zeroVec, zeroVec, zeroVec, zeroVec, zeroVec, coefdU, zeroVec, zeroVec, x};
        //double[][] coef = {k0, k1, zeoVec, x};

        double[] disp = new double[6];
        double[] force = new double[6];
        double presentforce = 0.0;


        int maxIteration = 10000;
        int iter = 0;
        NewIterSSolver solver = new NewIterSSolver();
        int idx = 1;
        while(idx < 6) {
            presentforce += 0.2;
            fs.assembleFlux(new double[]{0.2}, BVPType.COMMON, RHS, 2);
            iter = 0;
            while (iter < maxIteration) {
                //1. 计算余量
                A0.clear();
                fs.assembleStiff(coef, BVPType.NONLINEAR_1, A0);
                double[] R = MVO.add(RHS, -1.0, A0.mul(x));
                //2. 组装切线矩阵
                A.clear();
                fs.assembleStiff(coef, BVPType.NONLINEAR_2, A);

                //3. 施加边界条件
                fs.applyBC_MBN(A, RHS, Direct.All, 0.0, 1);

                //4. 求解增量dx
                solver.setMatrix(A);
                solver.PCGSSSOR(R, dx, 1.5, 0);

                //5. 计算整体解
                Arrays.setAll(x, i -> x[i] + dx[i]);

                //6. 检验收敛性
                if (MVO.L2Norm(dx) < 1.0e-10 * MVO.L2Norm(x) && MVO.L2Norm(R) < 1.0e-10 * MVO.L2Norm(RHS)) {
                    break;
                }
                iter++;
            }
            force[idx] = presentforce;
            disp[idx] = x[fs.GetNdof() - 1];
            System.out.println(force[idx] + "\t\t" + disp[idx]);
            idx++;
        }
        Tecplot.LineXY("LoadVsDisp2.dat", force, new double[][]{disp});
//        System.out.println(Arrays.toString(dx));
//        System.out.println(Arrays.toString(x));


    }
}
